"""
阶段1：基础截图功能
阶段2：自动滚动与多图拼接
阶段2增强：添加专用配置，支持登录状态保存
功能：初始化 Chrome WebDriver 并实现完整页面滚动截图
"""

import sys
import io

# 设置标准输出编码为UTF-8（Windows兼容性）
if sys.platform == 'win32':
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.common.exceptions import TimeoutException, WebDriverException
from webdriver_manager.chrome import ChromeDriverManager
from PIL import Image
from io import BytesIO
from datetime import datetime
import time
import os
import argparse

def basic_screenshot(url, save_path=None):
    """
    基础单屏截图功能
    
    参数:
        url (str): 目标网页URL
        save_path (str): 保存路径，默认使用时间戳命名
        
    返回:
        tuple: (final_url, save_path) 或 (None, None) 如果失败
    """
    driver = None
    
    try:
        print("=" * 60)
        print("阶段1：环境搭建与基础截图测试")
        print("=" * 60)
        
        # 步骤 1: 配置 Chrome 选项
        print("\n[1/5] 配置 Chrome 浏览器选项...")
        chrome_options = Options()
        chrome_options.add_argument('--headless=new')  # 无头模式
        chrome_options.add_argument('--disable-gpu')
        chrome_options.add_argument('--no-sandbox')
        chrome_options.add_argument('--disable-dev-shm-usage')
        chrome_options.add_argument('--window-size=1920,1080')
        chrome_options.add_argument('--disable-blink-features=AutomationControlled')
        print("  ✓ Chrome 选项配置完成")
        
        # 步骤 2: 初始化 WebDriver
        print("\n[2/5] 初始化 Chrome WebDriver...")
        
        # 尝试使用本地 chromedriver.exe
        local_chromedriver = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'chromedriver.exe')
        
        if os.path.exists(local_chromedriver):
            print(f"  ✓ 使用本地 ChromeDriver")
            service = Service(local_chromedriver)
        else:
            print("  ⚠ 未找到本地 ChromeDriver，尝试自动下载...")
            service = Service(ChromeDriverManager().install())
        
        driver = webdriver.Chrome(service=service, options=chrome_options)
        driver.implicitly_wait(10)
        driver.set_page_load_timeout(30)
        print("  ✓ WebDriver 初始化成功")
        
        # 步骤 3: 加载页面
        print(f"\n[3/5] 正在加载页面: {url}")
        driver.get(url)
        print("  ✓ 页面加载完成")
        
        # 步骤 4: 获取页面信息
        print("\n[4/5] 获取页面信息...")
        final_url = driver.current_url
        page_title = driver.title
        viewport_width = driver.execute_script("return window.innerWidth")
        viewport_height = driver.execute_script("return window.innerHeight")
        
        print(f"  页面标题: {page_title}")
        print(f"  最终URL: {final_url}")
        print(f"  视口尺寸: {viewport_width}x{viewport_height}")
        
        # 步骤 5: 截图并保存
        print("\n[5/5] 截取并保存截图...")
        
        # 自动生成文件名
        if save_path is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            save_path = f"basic_screenshot_{timestamp}.png"
        
        # 确保保存目录存在
        save_dir = os.path.dirname(save_path) if os.path.dirname(save_path) else "."
        if save_dir != "." and not os.path.exists(save_dir):
            os.makedirs(save_dir)
        
        # 截图
        driver.save_screenshot(save_path)
        
        # 获取文件大小
        file_size = os.path.getsize(save_path) / 1024  # KB
        
        print(f"  ✓ 截图已保存: {save_path}")
        print(f"  ✓ 文件大小: {file_size:.2f} KB")
        
        print("\n" + "=" * 60)
        print("✓ 阶段1测试成功！所有功能正常工作。")
        print("=" * 60)
        
        return final_url, save_path
        
    except TimeoutException as e:
        print(f"\n✗ 错误: 页面加载超时")
        print(f"  详情: {e}")
        return None, None
        
    except WebDriverException as e:
        print(f"\n✗ 错误: WebDriver 异常")
        print(f"  详情: {e}")
        return None, None
        
    except Exception as e:
        print(f"\n✗ 错误: 发生未知错误")
        print(f"  详情: {e}")
        import traceback
        traceback.print_exc()
        return None, None
        
    finally:
        # 清理资源
        if driver:
            driver.quit()
            print("\n✓ 浏览器已关闭，资源已释放")


def capture_full_page_screenshot(url, save_path=None, scroll_pause_time=0.5, setup_mode=False):
    """
    捕获网页完整滚动截图（阶段2功能 + 配置保存）
    
    参数:
        url (str): 目标网页URL
        save_path (str): 保存路径，默认使用时间戳命名
        scroll_pause_time (float): 滚动后等待时间（秒），懒加载页面建议1-2秒
        setup_mode (bool): 登录模式，显示浏览器窗口方便登录
        
    返回:
        tuple: (final_url, save_path) 或 (None, None) 如果失败
    """
    driver = None
    
    try:
        print("=" * 60)
        if setup_mode:
            print("🔧 登录模式 - 可以看到浏览器窗口，方便登录网站")
        else:
            print("📸 截图模式 - 后台运行，快速截图")
        print("=" * 60)
        
        # 步骤 1: 配置 Chrome 选项（带专用配置目录）
        print("\n[1/6] 配置 Chrome 浏览器选项...")
        chrome_options = Options()
        
        # 专用配置目录 - 保存登录状态
        profile_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'chrome_profile')
        chrome_options.add_argument(f'--user-data-dir={profile_dir}')
        chrome_options.add_argument('--profile-directory=ScreenshotProfile')
        
        # Setup模式：显示浏览器；默认模式：无头模式
        if not setup_mode:
            chrome_options.add_argument('--headless=new')  # 无头模式
        
        chrome_options.add_argument('--disable-gpu')
        chrome_options.add_argument('--no-sandbox')
        chrome_options.add_argument('--disable-dev-shm-usage')
        chrome_options.add_argument('--window-size=1920,1080')
        chrome_options.add_argument('--disable-blink-features=AutomationControlled')
        
        if setup_mode:
            print("  ✓ 登录模式已启用（可见浏览器窗口）")
        else:
            print("  ✓ 截图模式已启用（后台运行）")
        print(f"  ✓ 配置目录: {profile_dir}")
        print("  ✓ Chrome 选项配置完成")
        
        # 步骤 2: 初始化 WebDriver
        print("\n[2/6] 初始化 Chrome WebDriver...")
        
        # 尝试使用本地 chromedriver.exe，如果不存在则使用 webdriver_manager
        local_chromedriver = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'chromedriver.exe')
        
        if os.path.exists(local_chromedriver):
            print(f"  ✓ 使用本地 ChromeDriver: {local_chromedriver}")
            service = Service(local_chromedriver)
        else:
            print("  ⚠ 未找到本地 ChromeDriver，尝试自动下载...")
            try:
                service = Service(ChromeDriverManager().install())
            except Exception as e:
                print(f"  ✗ 自动下载失败: {e}")
                print("\n💡 解决方法：")
                print("  1. 下载 ChromeDriver: https://googlechromelabs.github.io/chrome-for-testing/")
                print("  2. 或国内镜像: https://registry.npmmirror.com/binary.html?path=chromedriver/")
                print(f"  3. 将 chromedriver.exe 放到: {local_chromedriver}")
                raise
        
        driver = webdriver.Chrome(service=service, options=chrome_options)
        driver.implicitly_wait(10)
        driver.set_page_load_timeout(30)
        print("  ✓ WebDriver 初始化成功")
        
        # 步骤 3: 加载页面并获取尺寸
        print(f"\n[3/6] 正在加载页面: {url}")
        driver.get(url)
        time.sleep(3)  # 等待页面初始加载
        print("  ✓ 页面加载完成")
        
        # Setup模式：提示用户可以登录
        if setup_mode:
            print("\n" + "🔐" * 30)
            print("  🔐 登录提示：")
            print("  - 如果需要登录，请在浏览器窗口中登录")
            print("  - 登录完成后，登录状态会自动保存")
            print("  - 下次使用时会自动保持登录状态")
            print("  - 准备好后请按回车继续截图...")
            print("🔐" * 30)
            input("\n按回车键继续...")
            print("  ✓ 继续截图流程")
        
        print("\n[4/6] 获取页面尺寸...")
        # 回到顶部
        driver.execute_script("window.scrollTo(0, 0);")
        time.sleep(0.5)
        
        # 获取页面总高度（多属性取最大值，更可靠）
        total_height = driver.execute_script("""
            return Math.max(
                document.body.scrollHeight,
                document.body.offsetHeight,
                document.documentElement.clientHeight,
                document.documentElement.scrollHeight,
                document.documentElement.offsetHeight
            );
        """)
        
        # 获取视口尺寸 - 使用实际截图的尺寸
        viewport_width = driver.execute_script("return document.documentElement.clientWidth")
        
        # 先截一张图，获取实际截图高度
        test_screenshot = driver.get_screenshot_as_png()
        test_image = Image.open(BytesIO(test_screenshot))
        actual_screenshot_width, actual_screenshot_height = test_image.size
        
        # 检测固定/粘性元素的高度（排除全屏元素）
        sticky_header_height = driver.execute_script("""
            let maxHeight = 0;
            const viewportHeight = window.innerHeight;
            const elements = document.querySelectorAll('*');
            
            elements.forEach(el => {
                const style = window.getComputedStyle(el);
                const position = style.position;
                
                if (position === 'fixed' || position === 'sticky') {
                    const rect = el.getBoundingClientRect();
                    
                    // 只考虑顶部的固定元素（top < 100px）且高度合理
                    // 排除全屏或接近全屏的元素（高度超过视口80%）
                    if (rect.top < 100 && rect.height > 0 && rect.height < viewportHeight * 0.8) {
                        maxHeight = Math.max(maxHeight, rect.bottom);
                    }
                }
            });
            
            return Math.floor(maxHeight);
        """)
        
        print(f"  页面总高度: {total_height}px")
        print(f"  浏览器视口宽度: {viewport_width}px")
        print(f"  实际截图尺寸: {actual_screenshot_width}x{actual_screenshot_height}px")
        if sticky_header_height > 0:
            print(f"  检测到粘性元素高度: {sticky_header_height}px")
            print(f"  将自动处理粘性元素重叠问题")
        print(f"  预计截图次数: {(total_height // (actual_screenshot_height - sticky_header_height)) + 1 if sticky_header_height > 0 else (total_height // actual_screenshot_height) + 1}")
        
        # 步骤 4: 滚动截图与拼接
        print("\n[5/6] 开始滚动截图与拼接...")
        # 创建空白长图 - 使用实际截图宽度
        stitched_image = Image.new('RGB', (actual_screenshot_width, total_height))
        
        # 计算滚动步长（考虑粘性元素），添加最小值保护
        scroll_step = actual_screenshot_height - sticky_header_height
        min_scroll_step = max(100, int(actual_screenshot_height * 0.2))  # 至少100px或视口的20%
        
        if scroll_step < min_scroll_step:
            print(f"  ⚠ 粘性元素过大({sticky_header_height}px)，使用最小滚动步长: {min_scroll_step}px")
            scroll_step = min_scroll_step
            sticky_header_height = actual_screenshot_height - scroll_step
        
        print(f"  滚动步长: {scroll_step}px")
        
        current_scroll_position = 0
        paste_position = 0
        screenshot_count = 0
        last_paste_position = -1  # 用于检测死循环
        stuck_count = 0  # 连续卡住的次数
        max_screenshots = 200  # 最大截图次数限制
        
        while paste_position < total_height and screenshot_count < max_screenshots:
            # 死循环检测
            if paste_position == last_paste_position:
                stuck_count += 1
                if stuck_count > 5:
                    print(f"  ⚠ 检测到死循环，强制终止截图")
                    break
            else:
                stuck_count = 0
                last_paste_position = paste_position
            
            # 滚动到当前位置
            driver.execute_script(f"window.scrollTo(0, {current_scroll_position});")
            time.sleep(scroll_pause_time)  # 等待页面渲染
            
            # 获取实际滚动位置（重要：浏览器可能无法滚动到目标位置）
            actual_scroll_y = driver.execute_script("return window.pageYOffset || document.documentElement.scrollTop;")
            
            # 截取当前视口
            screenshot = driver.get_screenshot_as_png()
            screenshot_image = Image.open(BytesIO(screenshot))
            screenshot_count += 1
            
            # 计算裁剪区域（优化后的逻辑）
            if screenshot_count == 1:
                # 第一张截图：完整保留
                crop_top = 0
                crop_height = actual_screenshot_height
            else:
                # 后续截图：需要考虑实际滚动位置和粘性元素
                # 截图内容对应页面范围：[actual_scroll_y, actual_scroll_y + actual_screenshot_height]
                # 期望开始粘贴的位置：paste_position
                
                # 基础裁剪：去掉顶部粘性元素
                crop_top = sticky_header_height
                
                # 检查是否有重叠（当浏览器无法滚动到期望位置时会发生）
                # 例如：页面500px，视口400px，第二次想滚动到400px但实际只能到100px
                if actual_scroll_y < paste_position:
                    # 截图内容的起始位置(actual_scroll_y)早于期望粘贴位置(paste_position)
                    # 说明截图顶部包含了已经粘贴过的内容，需要额外裁剪
                    overlap = paste_position - actual_scroll_y
                    crop_top = overlap + sticky_header_height
                
                # 计算可用高度（截图高度 - 裁剪的顶部）
                crop_height = actual_screenshot_height - crop_top
                
                # 确保不超过剩余需要截取的高度
                remaining_height = total_height - paste_position
                crop_height = min(crop_height, remaining_height)
            
            # 安全检查：确保裁剪参数有效
            if crop_height <= 0:
                print(f"  ⚠ 裁剪高度异常 ({crop_height}px)，停止截图")
                break
            
            # 裁剪图片
            screenshot_image = screenshot_image.crop((
                0,
                crop_top,
                actual_screenshot_width,
                crop_top + crop_height
            ))
            
            # 粘贴到长图
            stitched_image.paste(screenshot_image, (0, paste_position))
            
            # 更新位置
            paste_position += crop_height
            current_scroll_position += scroll_step
            
            # 显示进度
            progress = min(100, int(paste_position / total_height * 100))
            print(f"  截图进度: {progress}% (已截取 {screenshot_count} 张, 滚动: {actual_scroll_y}px, 裁剪: {crop_top}→{crop_top+crop_height}px, 粘贴: {paste_position}px)")
            
            # 如果已经完成，提前退出
            if paste_position >= total_height:
                break
        
        # 检查是否正常完成
        if screenshot_count >= max_screenshots:
            print(f"\n  ⚠ 警告：达到最大截图次数限制({max_screenshots}张)，截图可能不完整")
        elif stuck_count > 0:
            print(f"\n  ⚠ 检测到异常，提前终止")
        
        print(f"\n  ✓ 滚动截图完成，共截取 {screenshot_count} 张图片")
        print(f"  ✓ 图片拼接完成")
        
        # 步骤 5: 获取URL并保存
        print("\n[6/6] 保存截图...")
        final_url = driver.current_url
        page_title = driver.title
        
        # 自动生成文件名
        if save_path is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            save_path = f"fullpage_screenshot_{timestamp}.png"
        
        # 确保保存目录存在
        save_dir = os.path.dirname(save_path) if os.path.dirname(save_path) else "."
        if save_dir != "." and not os.path.exists(save_dir):
            os.makedirs(save_dir)
        
        # 保存截图
        stitched_image.save(save_path, 'PNG', optimize=True)
        
        # 获取文件大小
        file_size = os.path.getsize(save_path) / 1024  # KB
        
        print(f"  ✓ 截图已保存: {save_path}")
        print(f"  ✓ 文件大小: {file_size:.2f} KB")
        print(f"  ✓ 最终尺寸: {viewport_width}x{total_height}px")
        print(f"  ✓ 页面URL: {final_url}")
        print(f"  ✓ 页面标题: {page_title}")
        
        print("\n" + "=" * 60)
        print("✓ 阶段2测试成功！滚动截图与拼接功能正常工作。")
        print("=" * 60)
        
        return final_url, save_path
        
    except TimeoutException as e:
        print(f"\n✗ 错误: 页面加载超时")
        print(f"  详情: {e}")
        return None, None
        
    except WebDriverException as e:
        print(f"\n✗ 错误: WebDriver 异常")
        print(f"  详情: {e}")
        return None, None
        
    except Exception as e:
        print(f"\n✗ 错误: 发生未知错误")
        print(f"  详情: {e}")
        import traceback
        traceback.print_exc()
        return None, None
        
    finally:
        # 清理资源
        if driver:
            driver.quit()
            print("\n✓ 浏览器已关闭，资源已释放")


if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser(
        description='网页整页截图工具 - 支持登录状态保存',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
使用示例:
  # 完整滚动长截图（默认）
  python basic_screenshot.py --url https://example.com
  
  # 基础单屏截图
  python basic_screenshot.py --url https://example.com --basic
  
  # 登录模式（第一次使用，用于登录网站）
  python basic_screenshot.py --setup --url https://example.com
  
  # 交互式选择（无参数运行）
  python basic_screenshot.py
  
  # 自定义保存路径
  python basic_screenshot.py --url https://example.com --output screenshot.png
        '''
    )
    
    parser.add_argument('--setup', action='store_true', 
                       help='登录模式：显示浏览器窗口，方便登录网站（登录状态会保存）')
    parser.add_argument('--url', type=str, 
                       help='要截图的网页URL')
    parser.add_argument('--output', type=str, 
                       help='截图保存路径（可选）')
    parser.add_argument('--scroll-time', type=float, default=0.5,
                       help='滚动等待时间（秒，默认0.5）')
    parser.add_argument('--basic', action='store_true',
                       help='基础单屏截图模式（不滚动，只截取可见区域）')
    
    args = parser.parse_args()
    
    print("\n" + "=" * 60)
    print("🌐 网页整页截图工具")
    if args.setup:
        print("模式: 🔧 登录模式（可见浏览器窗口）")
    else:
        print("模式: 📸 截图模式（后台运行）")
    
    if args.basic:
        print("类型: 📄 基础单屏截图")
    else:
        print("类型: 📜 完整滚动长截图")
    print("=" * 60)
    
    # 如果提供了URL参数，直接使用
    if args.url:
        test_url = args.url
        print(f"\n目标URL: {test_url}")
        
        # 根据参数选择截图方式
        if args.basic:
            # 基础单屏截图
            final_url, save_path = basic_screenshot(test_url, args.output)
        else:
            # 完整滚动截图
            final_url, save_path = capture_full_page_screenshot(
                url=test_url,
                save_path=args.output,
                scroll_pause_time=args.scroll_time,
                setup_mode=args.setup
            )
        
    else:
        # 没有URL参数，交互式输入
        if args.setup:
            print("\n💡 提示：登录模式用于第一次登录网站，登录信息会保存")
            print("   下次使用截图模式时，会自动保持登录状态")
        
        # 选择截图类型（如果没有通过参数指定）
        if not args.basic:
            print("\n请选择截图类型：")
            print("1. 📜 完整滚动长截图（默认，适合长页面）")
            print("2. 📄 基础单屏截图（只截取可见区域）")
            
            screenshot_type = input("\n请输入选项 (1/2，默认为1): ").strip() or "1"
            use_basic = (screenshot_type == "2")
        else:
            use_basic = True
        
        print("\n请选择URL来源：")
        print("1. 输入自定义URL")
        print("2. 使用测试URL（Microsoft Edge 搜索页面）")
        
        choice = input("\n请输入选项 (1/2，默认为1): ").strip() or "1"
        
        if choice == "1":
            test_url = input("\n请输入要截图的URL: ").strip()
            if not test_url:
                print("❌ 错误：URL不能为空")
                sys.exit(1)
        else:
            test_url = "https://www.bing.com/search?q=screenshot+tool&form=QBLH"
            print(f"\n使用测试URL: Microsoft Edge 搜索结果页面")
        
        # 执行截图
        print("\n开始截图...")
        if use_basic:
            final_url, save_path = basic_screenshot(test_url, args.output)
        else:
            final_url, save_path = capture_full_page_screenshot(
                url=test_url,
                save_path=args.output,
                scroll_pause_time=args.scroll_time,
                setup_mode=args.setup
            )
    
    # 输出结果
    if save_path:
        print(f"\n" + "=" * 60)
        print("✓✓✓ 截图完成！")
        print(f"✓ 截图文件: {save_path}")
        print(f"✓ 最终URL: {final_url}")
        
        if args.setup:
            print("\n💡 下次使用：")
            print("   python basic_screenshot.py --url " + test_url)
            print("   （会自动保持登录状态，无需再登录）")
        
        print("=" * 60)
    else:
        print("\n" + "=" * 60)
        print("✗✗✗ 截图失败！请检查上方错误信息")
        print("=" * 60)

